import xml.etree.ElementTree as ET
import os

# 设置类别编号
classes = {'person': 0, 'bird': 1}


# 定义处理XML文件的函数
def process_xml(xml_file):
    # 加载XML文件
    tree = ET.parse(xml_file)
    root = tree.getroot()

    # 遍历XML文件中的所有object标签，提取出标注信息
    labels = []
    for obj in root.findall('object'):
        name = obj.find('name').text
        if name in classes:
            label = classes[name]
            bndbox = obj.find('bndbox')
            xmin = int(bndbox.find('xmin').text)
            ymin = int(bndbox.find('ymin').text)
            xmax = int(bndbox.find('xmax').text)
            ymax = int(bndbox.find('ymax').text)
            labels.append([label, xmin, ymin, xmax, ymax])

    return labels


# 处理所有的XML文件，并将标注信息导出成labels文件
xml_folder = '/outputs'
labels_file = '/xique_labels'

with open(labels_file, 'w') as f:
    for xml_file in os.listdir(xml_folder):
        if xml_file.endswith('.xml'):
            xml_path = os.path.join(xml_folder, xml_file)
            labels = process_xml(xml_path)
            for label in labels:
                f.write('{} {} {} {} {}\n'.format(label[0], label[1], label[2], label[3], label[4]))
